package com.dragon.utils;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeoutException;

@Slf4j
public class DashScopeUtils {
    public static Flux<String> runFlux(GenerationParam param) {
        Generation gen = new Generation();

        // 从SDK获取 Flowable 流。
        try {
            Flowable<GenerationResult> resultFlowable = gen.streamCall(param);
            // 使用 reactor-adapter 将 Flowable 转换为 Flux。
            return Flux.from(resultFlowable)
                    .timeout(Duration.ofSeconds(30)) // [AI Assistant] 建议增加超时，防止服务长时间等待
                    // 使用 flatMap 处理每个数据块。
                    .flatMap(generationResult -> {
                        String content = generationResult.getOutput().getChoices().get(0).getMessage().getContent();

                        if(content !=null && !content.isEmpty()){
                            log.info(content);
                        }
                        // 直接返回 content 即可
                        return Flux.just(Objects.requireNonNull(content));
                    })
                    .filter(str -> str != null && !str.isEmpty())
                    .onErrorResume(e -> {
                        if (e instanceof TimeoutException) {
                            log.error("调用DashScope API超时");
                            return Flux.just("AI分析超时，请稍后再试。");
                        }
                        log.error("调用DashScope API时发生错误", e);
                        return Flux.error(new RuntimeException("AI服务暂时不可用", e));
                    });
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static CompletableFuture<String> runNoFlux(GenerationParam param) {
        Generation gen = new Generation();
        // 创建一个 CompletableFuture 作为未来结果的容器
        CompletableFuture<String> future = new CompletableFuture<>();
        // 创建一个 StringBuilder 来拼接流式返回的文本块
        StringBuilder fullContent = new StringBuilder();
        try {
            gen.streamCall(param, new ResultCallback<>() {
                @Override
                public void onEvent(GenerationResult message) {
                    // 当收到流式返回的每一部分数据时，此方法被调用
                    String contentPiece = message.getOutput().getChoices().get(0).getMessage().getContent();
                    if (contentPiece != null) {
                        fullContent.append(contentPiece);
                    }
                }

                @Override
                public void onComplete() {
                    // 当所有数据都接收完毕时，此方法被调用
                    log.info("流执行完毕.");
                    future.complete(fullContent.toString());
                }

                @Override
                public void onError(Exception e) {
                    // 当发生错误时，此方法被调用
                    log.error("异步发生错误: " + e.getMessage());
                    future.completeExceptionally(e);
                }
            });
            return future;
        } catch (NoApiKeyException | InputRequiredException e) {
            log.error("发生错误，API或输入不合法: {}", e.getMessage());
        }
        return future;
    }
}
